package com.zhang.third.day08;

import org.apache.flink.api.common.typeutils.base.StringSerializer;
import org.apache.flink.api.common.typeutils.base.VoidSerializer;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.TwoPhaseCommitSinkFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;

import java.io.BufferedWriter;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Random;

/**
 * @title: 自定义两阶段提交
 * @author: zhang
 * @date: 2022/4/13 07:00
 */
public class Example2 {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        env.enableCheckpointing(10 * 1000L);

        env
                .addSource(new SourceFunction<Long>() {
                    private boolean running = true;
                    private Random random = new Random();

                    @Override
                    public void run(SourceContext<Long> ctx) throws Exception {
                        while (running) {
                            ctx.collect(random.nextLong());
                            Thread.sleep(1000l);
                        }
                    }

                    @Override
                    public void cancel() {
                        running = false;
                    }
                })
                .addSink(new TransactionalFileSink());


        env.execute();
    }

    public static class TransactionalFileSink extends TwoPhaseCommitSinkFunction<Long, String, Void> {
        private BufferedWriter transactionWriter;

        public TransactionalFileSink() {
            super(StringSerializer.INSTANCE, VoidSerializer.INSTANCE);
        }

        // 开始一个事务，并返回事务的标识符
        @Override
        protected String beginTransaction() throws Exception {
            long timeNow = System.currentTimeMillis();
            long taskIdx = getRuntimeContext().getIndexOfThisSubtask();
            String transactionFileName = timeNow + "-" + taskIdx;
            Path tempFilePath = Paths.get("/Users/apple/IdeaProjects/flink_1.13/src/main/resources/temp/" + transactionFileName);
            Files.createFile(tempFilePath);
            this.transactionWriter = Files.newBufferedWriter(tempFilePath);
            System.out.println("begin transaction");
            return transactionFileName;
        }

        @Override
        protected void invoke(String transaction, Long value, Context context) throws Exception {
            transactionWriter.write(value.toString());
            transactionWriter.write("\n");
        }

        @Override
        protected void preCommit(String transaction) throws Exception {
            transactionWriter.flush();
            transactionWriter.close();
        }

        @Override
        protected void commit(String transactionFileName) {
            Path tempFilePath = Paths.get("/Users/apple/IdeaProjects/flink_1.13/src/main/resources/temp/" + transactionFileName);
            if (Files.exists(tempFilePath)) {
                try {
                    Path commitFilePath = Paths.get("/Users/apple/IdeaProjects/flink_1.13/src/main/resources/commit/" + transactionFileName);
                    Files.move(tempFilePath, commitFilePath);
                    System.out.println("commit complete");
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }

        @Override
        protected void abort(String transactionFileName) {
            Path tempFilePath = Paths.get("/Users/apple/IdeaProjects/flink_1.13/src/main/resources/temp/" + transactionFileName);
            if (Files.exists(tempFilePath)){
                try {
                    Files.delete(tempFilePath);
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }

        }
    }
}
